/*
 * Javlov - a Java toolkit for reinforcement learning with multi-agent support.
 * 
 * Copyright (c) 2009 Matthijs Snel
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import net.javlov.util.ArrayUtil;

//TODO take eligibility traces out, this is something algorithmic
//TODO take generic parameterization out

@SuppressWarnings("unchecked")
public abstract class TabularQFunction implements QFunction {

	private static final long serialVersionUID = -1678305368586324288L;

	/**
	 * Use this factory method if the size of the state space is unknown. Will result in slower
	 * performance than when using {@link #getInstance(int, int)}.
	 * 
	 * If the size of the state space is unknown because the space is continuous, you probably
	 * shouldn't be using this class, but one of the classes that implement function approximators /
	 * generalisations over states for the value function (yet to be implemented).
	 * 
	 * The use case for this method is one in which the state space is discrete and of 
	 * unknown size (e.g. because the environment is dynamic and hence the size of the state space
	 * is hard to predict), yet finite and small enough for it to make sense to be stored in a
	 * table (as a rough indication I'd say less than a million states).
	 * 
	 * @param actionSpaceSize the (upper bound on the) size of the action space.
	 * @return Tabular Q-Value function for one-step Q-learning with initial values 0.
	 */
	public static TabularQFunction getInstance(int actionSpaceSize) {
		return getInstance(actionSpaceSize, 0.0);
	}
	
	/**
	 * Same as {@link #getInstance(int, double)} but with eligibility traces, where the
	 * parameter {@code traceDecay} is the eligibility trace decay parameter, a single
	 * parameter to represent {@code gamma*lambda}.
	 * 
	 * NOTE: the returned Q-value function updates values and eligibility traces as
	 * 
	 * {@code 
	 * Q(s,a) = Q(s,a) + alpha*delta*e(s,a)
	 * e(s,a) = gamma*lambda*e(s,a)
	 * }
	 * 
	 * and as such implement "naive Q(lambda)". It is not suitable for Watkins's or 
	 * Peng's Q(lambda) algorithms.
	 * 
	 * NOTE2: AT THE MOMENT THIS IMPLEMENTATION WILL CRASH WHEN USED WITH Q-LEARNING.
	 * 
	 * @param actionSpaceSize the (upper bound on the) size of the action space.
	 * @param traceDecay eligibility trace decay parameter. Single parameter to
	 * represent {@code gamma*lambda}.
	 * @return Tabular Q-Value function for Q-learning with eligibility traces where 
	 * {@code lambda = traceDecay / gamma} and initial values {@code initValue}.
	 */
	public static TabularQFunction getInstance(int actionSpaceSize, double traceDecay) {
		if ( traceDecay < 0.001 )
			return new UnknownStateSpaceFunction(actionSpaceSize);
		return new TracedUnknownStateSpaceFunction(actionSpaceSize, traceDecay);
	}
	
	/**
	 * Use this factory method if the size of the state space is known. Will result in faster
	 * performance than when using {@link #getInstance(int)}.
	 * 
	 * When using this method, the value function constructs an array of the size of the
	 * state space size, and uses the ID of the states to index into the array. Therefore, the
	 * following points of caution should be observed:
	 * 
	 * 1. The provided state space size can not be exceeded (i.e. the actual state space size
	 * can not be larger than the provided one), or errors might start occurring when using 
	 * the value function.
	 * 
	 * 2. States don't change; more particularly, the state ID uniquely identifies a state. For example,
	 * in a static grid world, a state with a particular ID will always be the same (position
	 * in grid, etc).
	 * 
	 * @param stateSpaceSize the (upper bound on the) size of the state space.
	 * @param actionSpaceSize the (upper bound on the) size of the action space.
	 * @return Tabular Q-Value function for one-step Q-learning with initial values 0.
	 */
	public static TabularQFunction getInstance(int stateSpaceSize, int actionSpaceSize) {
		return new KnownStateSpaceFunction(stateSpaceSize, actionSpaceSize);
	}
	
	/**
	 * Same as {@link #getInstance(int, int, double)} but with eligibility traces, where the
	 * parameter {@code traceDecay} is the eligibility trace decay parameter, a single
	 * parameter to represent {@code gamma*lambda}.
	 * 
	 * NOTE: the returned Q-value function updates values and eligibility traces as
	 * 
	 * {@code 
	 * Q(s,a) = Q(s,a) + alpha*delta*e(s,a)
	 * e(s,a) = gamma*lambda*e(s,a)
	 * }
	 * 
	 * and as such implement "naive Q(lambda)". It is not suitable for Watkins's or 
	 * Peng's Q(lambda) algorithms.
	 * 
	 * NOTE2: AT THE MOMENT THIS IMPLEMENTATION WILL CRASH WHEN USED WITH Q-LEARNING.
	 * 
	 * @param stateSpaceSize the (upper bound on the) size of the state space.
	 * @param actionSpaceSize the (upper bound on the) size of the action space.
	 * @param initValue the initial value for the Q-values, used to initialise the table.
	 * @param traceDecay eligibility trace decay parameter. Single parameter to
	 * represent {@code gamma*lambda}.
	 * @return Tabular Q-Value function for Q-learning with eligibility traces where 
	 * {@code lambda = traceDecay / gamma} and initial values {@code initValue}.
	 */
	public static TabularQFunction getInstance(int stateSpaceSize, int actionSpaceSize, 
			double traceDecay) {
		if ( traceDecay < 0.001 )
			return new KnownStateSpaceFunction(stateSpaceSize, actionSpaceSize);
		return new TracedKnownStateSpaceFunction(stateSpaceSize, actionSpaceSize, traceDecay);
	}
	
	/**
	 * gamma*lambda for eligibility traces.
	 */
	protected double traceDecay;
	
	protected double initVal;
	
	protected int actionSpaceSize;
	
	/**
	 * Default constructor.
	 */
	protected TabularQFunction() {}
	
	/**
	 * @inheritDoc
	 */
	@Override
	public double getTraceDecay() {
		return traceDecay;
	}

	/**
	 * @inheritDoc
	 */
	@Override
	public void setTraceDecay(double decay) {
		traceDecay = decay;			
	}
	
	/**
	 * Initialises the Q-values to the intial values specified in the constructor.
	 * Override if they should be intialised differently.
	 * @return
	 */
	protected double[] getInitValues(State s) {
		double[] vals = new double[actionSpaceSize];
		for ( int i = 0; i < actionSpaceSize; i++ )
			vals[i] = initVal;
		return vals; 
	}
	
	public abstract int size();
	
	public abstract Set<State> getStates();
	
	public abstract void put(State s, double[] vals);
	
	public abstract double[] queryValues(State s);
	
	public static class UnknownStateSpaceFunction extends TabularQFunction {

		private static final long serialVersionUID = 8186346862373024852L;

		//Performance test showed Trove only outperforms Java HashMap if state space larger than
		//around 250,000 states or so; for smaller spaces, it is actually slower. On the assumption
		//that most uses will involve less states, using HashMap.
		protected Map<State, double[]> table;
		
		protected State lastState, prevLastState;
		protected double[] lastVals, prevLastVals;
		protected int lastActionID, prevLastActionID;
		
		protected UnknownStateSpaceFunction(int actionSpaceSize) {
			table = new HashMap<State, double[]>(5000);
			this.actionSpaceSize = actionSpaceSize;
		}
		
		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State, Action)
		 */
		@Override
		public double getValue(State s, Option opt) {
			updateMemory(s, opt);
			if ( lastVals == null )
				return getInitValues(s)[lastActionID];
			return lastVals[lastActionID];
		}

		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State)
		 */
		@Override
		public double[] getValues(State s) {
			updateMemory(s, null);
			if ( lastVals == null )
				return getInitValues(s);
			return lastVals;
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void setValue(State s, Option opt, double value) {
			double[] vals = table.get(s);
			if ( vals == null ) {
				vals = getInitValues(s);
				table.put(s, vals);
			}
			vals[opt.getID()] = value;
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void setLastOption(Option o) {
			lastActionID = o.getID();
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void update(State s, Option opt, double delta) {
			update(s, table.get(s), opt.getID(), delta);
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void update(double delta) {
			update(lastState, lastVals, lastActionID, delta);
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public void updatePrevious(double delta) {
			update(prevLastState, prevLastVals, prevLastActionID, delta);
		}
				
		protected void update(State s, double[] vals, int actID, double delta) {
			if ( vals == null ) {
				vals = getInitValues(s);
				table.put(s, vals);
			}
			vals[actID] += delta;
		}
		
		/**
		 * Updates the function's "memory" of previous states, actions, and retrieved
		 * (s,a) values, in order to be able to implement methods {@link #update(double)}
		 * and {@link #updatePrevious(double)}.
		 * 
		 * @param s the state currently being queried
		 * @param act the action currently being queried - can be null
		 */
		protected void updateMemory(State s, Option opt) {
			prevLastState = lastState;
			lastState = s;
			prevLastVals = lastVals;
			lastVals = table.get(s);
			prevLastActionID = lastActionID;
			if ( opt == null )
				lastActionID = 0;
			else
				lastActionID = opt.getID();
		}
		
		public void init() {
			reset();
			table.clear();
		}
		
		public void reset() {
			prevLastState = lastState = null;
			prevLastVals = lastVals = null;
			prevLastActionID = lastActionID = 0;
		}
		
		public String toString() {
			String s = "";
			for ( State st : table.keySet() ) {
				s += st.toString() + ":" + Arrays.toString(table.get(st)) + "\n";
			}
			return s;
		}
		
		public int size() {
			return table.size() * actionSpaceSize;
		}

		@Override
		public void setInitValues(double initValue) {
			initVal = initValue;			
		}
		
		@Override
		public Set<State> getStates() {
			return table.keySet();
		}
		
		@Override
		public void put(State s, double[] vals) {
			table.put(s, vals);
		}

		@Override
		public double[] queryValues(State s) {
			return table.get(s);
		}

	}
	
	public static class TracedUnknownStateSpaceFunction extends UnknownStateSpaceFunction {
		
		private static final long serialVersionUID = -823571748167107323L;
		
		/**
		 * Table of traces for each state-action pair. Use an array to also
		 * store the value of the pair while it's "tracing"; this avoids having to
		 * retrieve the value from the value table each time the pair is updated.
		 * 
		 * First entry in the array is trace, second is value.
		 */
		protected Map<State, double[][]>traceTable;
		
		protected TracedUnknownStateSpaceFunction(int actionSpaceSize, double traceDecay) {
			super(actionSpaceSize);
			traceTable = new HashMap<State, double[][]>(100);
			setTraceDecay(traceDecay);
		}
		
		@Override
		protected void update(State s, double[] vals, int actID, double delta) {
			//check if the (s,a) pair to update is already in the trace table.
			double eAndQ[][] = traceTable.get(s);
			if ( eAndQ == null ) { //state has no actions with traces
				eAndQ = new double[2][actionSpaceSize];
				
				if ( vals == null ) { //first time state is encountered
					vals = getInitValues(s);
					table.put(s, vals); //store init values in value table
				}
				eAndQ[1] = vals; //store a pointer to the values in the trace table
				traceTable.put(s, eAndQ);
			}
			//replace the trace of the action that we're currently updating, and set
			//traces of other actions to 0 (recommended by Sutton & Barto)
			for ( int i = 0; i < actionSpaceSize; i++ )
				eAndQ[0][i] = (i == actID ? 1 : 0);
			
			//now update all (s,a) pairs that have traces above threshold
			updateAll(delta);
		}
		
		/**
		 * Update all (s,a) pairs with trace e(s,a) > 0.001.
		 * 
		 * @param delta the delta to update with.
		 */
		protected void updateAll(double delta) {
			double traceAndVals[][];
			boolean remove = true; //flag entries that need removal (all action traces below threshold)
			List<Map.Entry<State, double[][]>> toRemove = new ArrayList<Map.Entry<State, double[][]>>();
			Set<Map.Entry<State, double[][]>> traceSet = traceTable.entrySet();
			for ( Map.Entry<State, double[][]> traceEntry : traceSet ) {
				traceAndVals = traceEntry.getValue();
				
				for ( int i = 0; i < actionSpaceSize; i++ ) {
					//update value: array in trace table is same pointer as in value table
					//Q(s,a) = Q(s,a) + alpha*delta*e(s,a)
					traceAndVals[1][i] += delta*traceAndVals[0][i];
					//decay trace: e(s,a) = gamma*lambda*e(s,a)
					if ( (traceAndVals[0][i] *= traceDecay) > 0.001 )
						remove = false;
				}
				if ( remove )
					toRemove.add(traceEntry);
				remove = true;
			}
			//remove here because we have to wait until iterator is finished
			traceSet.removeAll(toRemove);
		}
	}
	
	public static class KnownStateSpaceFunction extends TabularQFunction {
		
		private static final long serialVersionUID = 5529805948086768159L;
		
		/**
		 * Table that stores (s,a) values. 1st index = states, 2nd = actions.
		 */
		protected double[][] table;
		
		private int lastStateID, prevLastStateID;
		private int lastActionID, prevLastActionID;
		private int stateSpaceSize; 
		
		protected KnownStateSpaceFunction(int stateSpaceSize, int actionSpaceSize) {
			this.stateSpaceSize = stateSpaceSize;
			this.actionSpaceSize = actionSpaceSize;
			table = new double[stateSpaceSize][actionSpaceSize];
		}
		
		/**
		 * Initialises the table to 0 values. Override if it should be initialised differently.
		 * @param stateSpaceSize
		 * @param actionSpaceSize
		 * @return
		 */
		protected void initTable() {
			for ( int i = 0; i < stateSpaceSize; i++ )
				for ( int j = 0; j < actionSpaceSize; j++ )
					table[i][j] = initVal;
		}
		
		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State, Action)
		 */
		@Override
		public double getValue(State s, Option opt) {
			updateMemory(s, opt);
			return table[lastStateID][lastActionID];
		}

		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State)
		 */
		@Override
		public double[] getValues(State s) {
			updateMemory(s, null);
			return table[lastStateID];
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void setValue(State s, Option opt, double value) {
			table[s.getID()][opt.getID()] = value;	
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void setLastOption(Option opt) {
			lastActionID = opt.getID();
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void update(State s, Option opt, double delta) {
			update(s.getID(), opt.getID(), delta);			
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void update(double delta) {
			update(lastStateID, lastActionID, delta);
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public void updatePrevious(double delta) {
			update(prevLastStateID, prevLastActionID, delta);
		}
		
		protected void update(int stateID, int actionID, double delta) {
			table[stateID][actionID] += delta;
		}
		
		/**
		 * Updates the function's "memory" of previous states, actions, and retrieved
		 * (s,a) values, in order to be able to implement methods {@link #update(double)}
		 * and {@link #updatePrevious(double)}.
		 * 
		 * @param s the state currently being queried
		 * @param act the action currently being queried - can be null
		 */
		protected void updateMemory(State s, Option opt) {
			prevLastStateID = lastStateID;
			lastStateID = s.getID();
			prevLastActionID = lastActionID;
			if ( opt == null )
				lastActionID = 0;
			else
				lastActionID = opt.getID();
		}
		
		public void init() {
			reset();
			table = new double[table.length][table[0].length];
		}
		
		public void reset() {
			prevLastStateID = lastStateID = 0;
			prevLastActionID = lastActionID = 0;
		}
		
		public void setInitValues(double initValue) {
			initVal = initValue;
			initTable();
		}
		
		public String toString() {
			return ArrayUtil.arrayToString(table, 1);
		}

		@Override
		public int size() {
			return stateSpaceSize*actionSpaceSize;
		}
		
		@Override
		public Set<State> getStates() {
			//TODO
			return null;
		}
		
		@Override
		public void put(State s, double[] vals) {
			double[] copy = new double[vals.length];
			System.arraycopy(vals, 0, copy, 0, vals.length);
			table[s.getID()] = copy;
		}
		
		public double[] queryValues(State s) {
			return table[s.getID()];
		}
	}
	
	public static class TracedKnownStateSpaceFunction extends KnownStateSpaceFunction {
		
		private static final long serialVersionUID = 8512506806070140454L;

		/**
		 * Table with traces. Integer is the state id, array stores trace for each action.
		 */
		protected Map<Integer, double[]> traceTable;
		
		protected TracedKnownStateSpaceFunction(int stateSpaceSize, int actionSpaceSize,
				double traceDecay) {
			super(stateSpaceSize, actionSpaceSize);
			traceTable = new HashMap<Integer, double[]>(100);
			setTraceDecay(traceDecay);
		}
		
		@Override
		protected void update(int stateID, int actionID, double delta) {
			//check if the (s,a) pair to update is already in the trace table.
			double[] traces = traceTable.get(stateID);
			if ( traces == null ) {
				traces = new double[table[0].length];
				traceTable.put(stateID, traces);
			}
			//replace the trace of the action that we're currently updating, and set
			//traces of other actions to 0 (recommended by Sutton & Barto)
			for ( int i = 0; i < traces.length; i++ )
				traces[i] = (i == actionID ? 1 : 0);
			
			updateAll(delta);
		}
		
		/**
		 * Update all states that have traces above threshold 0.001.
		 * 
		 * @param delta the delta to update with.
		 */
		protected void updateAll(double delta) {
			double traces[];
			int stateID;
			boolean remove = true; //flag entries that need removal (all action traces below threshold)
			List<Map.Entry<Integer, double[]>> toRemove = new ArrayList<Map.Entry<Integer, double[]>>();
			Set<Map.Entry<Integer, double[]>> traceSet = traceTable.entrySet();
			for ( Map.Entry<Integer, double[]> traceEntry : traceSet ) {
				traces = traceEntry.getValue();
				stateID = traceEntry.getKey();
				
				for ( int i = 0; i < traces.length; i++ ) {
					double before = table[stateID][i];
					//update value: Q(s,a) = Q(s,a) + alpha*delta*e(s,a)
					table[stateID][i] += delta*traces[i];
					if ( Double.isNaN(table[stateID][i]) )
						System.err.println("Oops. " + before + "," + delta + "," + traces[i]);
					//decay trace: e(s,a) = gamma*lambda*e(s,a)
					if ( (traces[i] *= traceDecay) > 0.001 )
						remove = false;
				}
				if ( remove )
					toRemove.add(traceEntry);
				remove = true;
			}
			traceSet.removeAll(toRemove);
		}
	}
}
